#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <limits.h>
#include <assert.h>

void exclusive_scan(int *input, int length)
{
    if(length == 0 || length == 1)
        return;

    int old_val, new_val;

    old_val = input[0];
    input[0] = 0;
    for (int i = 1; i < length; i++)
    {
        new_val = input[i];
        input[i] = old_val + input[i-1];
        old_val = new_val;
    }
}

void convert_CSR_to_CSC(const int     m,
                          const int     n,
                          const int     nnz,
                          const int    *csrRowPtr,
                          const int    *csrColIdx,
                          const double *csrVal,
                          int          *cscRowIdx,
                          int          *cscColPtr,
                          double       *cscVal)
{
    // histogram in column pointer
    memset (cscColPtr, 0, sizeof(int) * (n+1));
    for (int i = 0; i < nnz; i++) {
        cscColPtr[csrColIdx[i]]++;
    }

    // prefix-sum scan to get the column pointer
    exclusive_scan(cscColPtr, n + 1);

    int *cscColIncr = (int *)malloc(sizeof(int) * (n+1));
    memcpy (cscColIncr, cscColPtr, sizeof(int) * (n+1));

    // insert nnz to csc
    for (int row = 0; row < m; row++)
    {
        for (int j = csrRowPtr[row]; j < csrRowPtr[row+1]; j++)
        {
            int col = csrColIdx[j];

            cscRowIdx[cscColIncr[col]] = row;
            cscVal[cscColIncr[col]] = csrVal[j];
            cscColIncr[col]++;
        }
    }

    free (cscColIncr);
}

int main(int argc, char* argv[]) {

    // 基本参数
    char *filename_matrix = argv[1];
    char *filename_matrix_cleaned = argv[2];
    int nrows = atoi(argv[3]);   
    int nnz = atoi(argv[4]);  
    printf("[Reading matrix from %s]\n", filename_matrix);

    // 读取矩阵的临时变量
    int rid, cid;
    double val;

    // CSR 
    int *row_ptr = new int[nrows+1];            
    int *col_idx = new int[nnz];    
    double *val_f64 = new double[nnz]; 
    double *row_cksum = new double[nrows];
    memset(row_ptr, 0, sizeof(int) * (nrows+1));
    memset(row_cksum, 0, sizeof(double) * nrows);
    double total_cksum = 0.0;

    // 开始读取
    FILE *fp = fopen(filename_matrix, "r");     
    for(int i=0;i<nnz;i++)
    {
        // 1-based index
        fscanf(fp, "%d %d %lf", &rid, &cid, &val);
        row_ptr[rid]++;
        col_idx[i] = cid-1;
        val_f64[i] = val;
        row_cksum[rid-1] += (val_f64[i] * col_idx[i]);
        total_cksum += val_f64[i];
    }
    fclose(fp);

    for(int i=1;i<=nrows;i++)
        row_ptr[i] += row_ptr[i-1];

    if(row_ptr[nrows] != nnz) printf("warning: row_ptr is wrong \n"); 

    printf("[Matrix load finish]\n");
    
    // 每一行内可能存在 col_idx 重复的情况，需要进行合并
    int * row_ptr_new = (int *) malloc(sizeof(int) * (nrows+1));
    int * col_idx_new = (int *) malloc(sizeof(int) * nnz);
    double * val_f64_new = (double *) malloc(sizeof(double) * nnz);
    int real_nnz = 0;

    for (int i = 0; i < nrows; i++) {
        int start = row_ptr[i];
        int end = row_ptr[i+1];
        // 排序
        for (int j = start; j < end; j++) {
            for (int k = j+1; k < end; k++) {
                if (col_idx[j] > col_idx[k]) {
                    int tmp = col_idx[j];
                    col_idx[j] = col_idx[k];
                    col_idx[k] = tmp;
                    double tmp2 = val_f64[j];
                    val_f64[j] = val_f64[k];
                    val_f64[k] = tmp2;
                }
            }
        }
        // 合并
        for (int j = start; j < end; j++) {
            if (j == start || col_idx[j] != col_idx[j-1]) {
                col_idx_new[real_nnz] = col_idx[j];
                val_f64_new[real_nnz] = val_f64[j];
                real_nnz++;
            } else {
                val_f64_new[real_nnz-1] += val_f64[j];
                // printf("[Find] row: %d, col: %d, val: %.6f\n", i, col_idx[j], val_f64[j]);
            }
        }
        // 检验
        double ck_sum = 0;
        for (int j = row_ptr_new[i]; j < real_nnz; j++) {
            ck_sum += (val_f64_new[j] * col_idx_new[j]);
        }
        if (fabs(ck_sum - row_cksum[i]) > 1e-6) {
            printf("[Error] row %d: %lf -> %lf, nnz: %d -> %d\n", i, row_cksum[i], ck_sum, row_ptr[i+1] - row_ptr[i], real_nnz - row_ptr_new[i]);
            for (int j = row_ptr_new[i]; j < real_nnz; j++) {
                printf("    col: %d, val: %.6f\n", col_idx_new[j], val_f64_new[j]);
            }
            break;
        }

        row_ptr_new[i+1] = real_nnz;
    }

    printf("[Matrix deduplicate finish]\n");
    printf("[original nnz: %d]\n", nnz);
    printf("[real     nnz: %d]\n", real_nnz);
    printf("[check sum = %.8f]\n", total_cksum);
    assert(row_ptr_new[nrows] == real_nnz);

    printf("[Matrix transpose finish]\n");
    printf(" ===> writing to binaries....\n");
    
    char target[200];
    int size;
    
    sprintf(target, "%s/Ai.bin\0", filename_matrix_cleaned);
    fp = fopen(target, "wb+");
    size = fwrite(row_ptr_new, sizeof(int), nrows+1, fp);
    printf("write %d into row_ptr at [%s]\n", size, target);
    fclose(fp);

    sprintf(target, "%s/Aj.bin\0", filename_matrix_cleaned);
    fp = fopen(target, "wb+");
    size = fwrite(col_idx_new, sizeof(int), real_nnz, fp);
    printf("write %d into col_idx at [%s]\n", size, target);
    fclose(fp);

    sprintf(target, "%s/Av.bin\0", filename_matrix_cleaned);
    fp = fopen(target, "wb+");
    size = fwrite(val_f64_new, sizeof(double), real_nnz, fp);
    printf("write %d into val at [%s]\n", size, target);
    fclose(fp);

    // int * trs_col_ptr = new int[nrows+1]; //行偏移量
    // int * trs_row_idx = new int[real_nnz];     //列索引
    // double * trs_val  = new double[real_nnz];  //非零元
    // convert_CSR_to_CSC(nrows, nrows, real_nnz, row_ptr_new, col_idx_new, val_f64_new, trs_row_idx, trs_col_ptr, trs_val);

    // sprintf(target, "%s.row_idx\0", filename_matrix_cleaned);
    // fp = fopen(target, "wb+");
    // size = fwrite(trs_row_idx, sizeof(int), real_nnz, fp);
    // printf("write %d into row_idx\n", size);
    // fclose(fp);

    // sprintf(target, "%s.col_ptr\0", filename_matrix_cleaned);
    // fp = fopen(target, "wb+");
    // size = fwrite(trs_col_ptr, sizeof(int), nrows+1, fp);
    // printf("write %d into row_idx\n", size);
    // fclose(fp);

    // sprintf(target, "%s.val_csc\0", filename_matrix_cleaned);
    // fp = fopen(target, "wb+");
    // size = fwrite(trs_val, sizeof(double), real_nnz, fp);
    // printf("write %d into row_idx\n", size);
    // fclose(fp);
    
    printf("\n");

    delete [] row_ptr;
    delete [] col_idx;
    delete [] val_f64;
    delete [] col_idx_new;
    delete [] val_f64_new;
    delete [] row_ptr_new;
    // delete [] trs_col_ptr;
    // delete [] trs_row_idx;
    // delete [] trs_val;

    return 0;
}

